#ifndef _MULTIGRID_H_
#define _MULTIGRID_H_

#include "Multigrid/PossionOP.h"
#include "Multigrid/BottomSolver.h"
#include "Multigrid/Restriction.h"
#include "Multigrid/Interpolation.h"
#include "RegularGrid/FuncFiller.h"
#include <vector>
#include <memory>
#include <cstring>

template <int Dim> class Multigrid;

template <int Dim>
class Multigrid{
public:
  template <class T>
  using Vector = std::vector<T>;

  using VPR = Vector<Restriction<Dim>*>;

  using VPI = Vector<Interpolation<Dim>*>;

  Multigrid() = default;

  Multigrid(const Vector<RectDomain<Dim> >& avDomain, const VPR&
            avpRestriction, const VPI& avpInterpolation);

  ~Multigrid();

  void setBCType(const char* aBCTypes);

  void setParam(const Real aweight, const int anumIter1, const int
                anumIter2, const int anumIter3, const Real areltol,
                const int amaxIter);

  const Vector<RectDomain<Dim> >& getvDomain() const;

  const char* getBCTypes() const;

  const Real getWeight() const;

  const Vec<int,3> getNumIter() const;

  const Real getReltol() const;

  const int getMaxIter() const;

  template <class TFunc>
  Real SolveVCycle(Tensor<Real,Dim>& phi, const Tensor<Real,Dim>& rhs,
  const TFunc& f);

  template <class TFunc>
  Real SolveFMVCycle(Tensor<Real,Dim>& phi, const Tensor<Real,Dim>& rhs,
  const TFunc& f);

  protected:
  void VCycle(int depth, Tensor<Real,Dim>& phi, const Tensor<Real,Dim>&
              rhs);

  void FMVCycle(int depth, Tensor<Real,Dim>& phi, const Tensor<Real,Dim>&
                            rhs);
  
protected:
  // Operator listed here
  Vector<RectDomain<Dim> > vDomain;
  Vector<GhostFiller<Dim> > vGhostFiller;
  Vector<PossionOP<Dim> > vPossionOP;
  VPR vpRestriction;
  VPI vpInterpolation;
  BottomSolver<Dim> bottomSolver;
  // Boundry condition
  char BCTypes[4];
  // Other parameter
  Real weight;
  int numIter1;
  int numIter2;
  int numIter3;
  Real reltol;
  int maxIter;
};

template <int Dim>
Multigrid<Dim>::Multigrid(const Vector<RectDomain<Dim> >& avDomain, const VPR&
                        avpRestriction, const VPI& avpInterpolation){
  vDomain = avDomain;
  vpRestriction = avpRestriction;
  vpInterpolation = avpInterpolation;
  const int size = avDomain.size();
  for (int i = 0 ; i < size ; i++){
    vGhostFiller.push_back(GhostFiller<Dim>{avDomain[i]});
    vPossionOP.push_back(PossionOP<Dim>{avDomain[i]});
  }  
  bottomSolver = BottomSolver<Dim>(avDomain[size-1]);
}

template <int Dim>
Multigrid<Dim>::~Multigrid(){}

template <int Dim>
void Multigrid<Dim>::setBCType(const char* aBCTypes){
  strcpy(BCTypes,aBCTypes);
}

template <int Dim>
void Multigrid<Dim>::setParam(const Real aweight, const int anumIter1,
                              const int anumIter2, const int anumIter3, const Real areltol,
                              const int amaxIter){
  weight = aweight;
  numIter1 = anumIter1;
  numIter2 = anumIter2;
  numIter3 = anumIter3;
  reltol = areltol;
  maxIter = amaxIter;
}

template <int Dim>
const std::vector<RectDomain<Dim> >& Multigrid<Dim>::getvDomain() const{
  return vDomain;
}

template <int Dim>
const char* Multigrid<Dim>::getBCTypes() const{
  return BCTypes;
}

template <int Dim>
const Real Multigrid<Dim>::getWeight() const{
  return weight;
}

template <int Dim>
const Real Multigrid<Dim>::getReltol() const{
  return reltol;
}

template <int Dim>
const Vec<int,3> Multigrid<Dim>::getNumIter() const{
  return Vec<int,3>{numIter1,numIter2,numIter3};
}

template <int Dim>
const int Multigrid<Dim>::getMaxIter() const{
  return maxIter;
}

template <int Dim>
void Multigrid<Dim>::VCycle(int depth, Tensor<Real,Dim>& phi, const Tensor<Real,Dim>& rhs){
  if (depth == (int)(vDomain.size()) - 1){
    bottomSolver.Solve(phi,rhs,BCTypes,weight,numIter3);
    return;
  }
  Tensor<Real,Dim> tmp(vDomain[depth].getGhostedBox());
  for (int i = 0 ; i < numIter1 ; i+=2){
    vGhostFiller[depth].fillAllSides(phi,BCTypes,0);
    vPossionOP[depth].relaxLaplacian(phi,rhs,tmp,weight);
    vGhostFiller[depth].fillAllSides(tmp,BCTypes,0);
    vPossionOP[depth].relaxLaplacian(tmp,rhs,phi,weight);
  }
  vGhostFiller[depth].fillAllSides(phi,BCTypes,0);
  Tensor<Real,Dim> rsd(vDomain[depth].getGhostedBox());
  vPossionOP[depth].computeResidul(phi,rhs,rsd);
  vGhostFiller[depth].fillAllSides(rsd,BCTypes,0);
  Tensor<Real,Dim> crsd(vDomain[depth+1].getGhostedBox());
  vpRestriction[depth]->applyRestrict(rsd,crsd);
  Tensor<Real,Dim> cphi(vDomain[depth+1].getGhostedBox());
  cphi = 0.0;
  VCycle(depth+1,cphi,crsd);
  vpInterpolation[depth]->applyInterpolate(cphi,phi);
  for (int i = 0 ; i < numIter2 ; i+=2){
    vGhostFiller[depth].fillAllSides(phi,BCTypes,0);
    vPossionOP[depth].relaxLaplacian(phi,rhs,tmp,weight);
    vGhostFiller[depth].fillAllSides(tmp,BCTypes,0);
    vPossionOP[depth].relaxLaplacian(tmp,rhs,phi,weight);
  }
  vGhostFiller[depth].fillAllSides(phi,BCTypes,0);
}

template <int Dim>
template <class TFunc>
Real Multigrid<Dim>::SolveVCycle(Tensor<Real,Dim>& phi, const Tensor<Real,Dim>& rhs,
                                 const TFunc& f){
  const RectDomain<Dim>& Domain = vDomain[0];
  Tensor<Real,Dim> rsd(Domain.getGhostedBox());
  vGhostFiller[0].fillAllSides(phi,BCTypes,f);
  vPossionOP[0].computeResidul(phi,rhs,rsd);
  vGhostFiller[0].fillAllSides(rsd,BCTypes,0);
  Real iRsd = vPossionOP[0].computeNorm(rsd);
  Real Rsd[2];
  Rsd[0] = Rsd[1] = iRsd;
  int count = 1;
  Real relRsd,RsdRatio;
  Tensor<Real,Dim> res(Domain.getGhostedBox());
  while(1){
    Rsd[0] = Rsd[1];
    res = 0.0;
    VCycle(0,res,rsd);
    phi = phi + res;
    vGhostFiller[0].fillAllSides(phi,BCTypes,f);
    vPossionOP[0].computeResidul(phi,rhs,rsd);
    vGhostFiller[0].fillAllSides(rsd,BCTypes,0);
    Rsd[1] = vPossionOP[0].computeNorm(rsd);
    relRsd = Rsd[1]/iRsd;
    RsdRatio = Rsd[1]/Rsd[0];
    std::cout << "Iter " << count << " , relrsd = " << relRsd <<\
      " , rsdratio = " << RsdRatio << std::endl;
    if (relRsd < reltol || count > maxIter)
      break;
    count++;
  }
  return relRsd;
}

template <int Dim>
void Multigrid<Dim>::FMVCycle(int depth, Tensor<Real,Dim>& phi, const Tensor<Real,Dim>&
                              rhs){
  if (depth != (int)(vDomain.size()) - 1){
    Tensor<Real,Dim> crhs(vDomain[depth+1].getGhostedBox());
    vpRestriction[depth]->applyRestrict(rhs,crhs);
    Tensor<Real,Dim> cphi(vDomain[depth+1].getGhostedBox());
    cphi = 0.0;
    FMVCycle(depth+1,cphi,crhs);
    vpInterpolation[depth]->applyInterpolate(cphi,phi);
  }
  VCycle(depth,phi,rhs);
}


template <int Dim>
template <class TFunc>
Real Multigrid<Dim>::SolveFMVCycle(Tensor<Real,Dim>& phi, const Tensor<Real,Dim>& rhs,
                                   const TFunc& f){
   const RectDomain<Dim>& Domain = vDomain[0];
   Tensor<Real,Dim> rsd(Domain.getGhostedBox());
   vGhostFiller[0].fillAllSides(phi,BCTypes,f);
   vPossionOP[0].computeResidul(phi,rhs,rsd);
   vGhostFiller[0].fillAllSides(rsd,BCTypes,0);
   Real iRsd = vPossionOP[0].computeNorm(rsd);
   Real Rsd[2];
   Rsd[0] = Rsd[1] = iRsd;
   int count = 1;
   Real relRsd;
   //Real RsdRatio;
   Tensor<Real,Dim> res(Domain.getGhostedBox());
   while(1){
     Rsd[0] = Rsd[1];
     res = 0.0;
     FMVCycle(0,res,rsd);
     phi = phi + res;
     vGhostFiller[0].fillAllSides(phi,BCTypes,f);
     vPossionOP[0].computeResidul(phi,rhs,rsd);
     vGhostFiller[0].fillAllSides(rsd,BCTypes,0);
     Rsd[1] = vPossionOP[0].computeNorm(rsd);
     relRsd = Rsd[1]/iRsd;
     //RsdRatio = Rsd[1]/Rsd[0];
     // std::cout << "Iter " << count << " , relrsd = " << relRsd <<       
     //   " , rsdratio = " << RsdRatio << std::endl;
     if (relRsd < reltol || count > maxIter)
       break;
     count++;
   }
   return relRsd;
}
#endif // _MULTIGRID_H_
